# 恋爱技巧提取
import os
import argparse
from tqdm import tqdm
from client import get_by_name
from joblib import Parallel, delayed

def distill_skill(transcript_path, output_path, model):
    """
    从转写文本中提取技巧并保存结果
    
    Parameters:
    transcript_path (str): 转写文本文件路径
    output_path (str): 输出文件路径
    model: 大模型客户端
    
    Returns:
    bool: 处理是否成功
    """
    # 读取转写文本
    try:
        with open(transcript_path, 'r', encoding='utf-8') as f:
            transcript = f.read()
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 使用大模型提取技巧
        PROMPT = """
        {transcript}

        以上是撩妹课程视频的字幕内容，帮我从中抽取有价值的撩妹技巧及实战案例。
        """
        
        response = model(PROMPT.format(transcript=transcript))
        
        # 将结果写入输出文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(response)
            
        return (transcript_path, True)
    except Exception as e:
        print(f"处理文件 {transcript_path} 时出错: {e}")
        return (transcript_path, False)

def process_single_file(transcript_file, input_dir, output_dir, model):
    """
    处理单个文件的包装函数，用于并行处理
    
    Parameters:
    transcript_file (str): 转写文本文件路径
    input_dir (str): 输入目录路径
    output_dir (str): 输出目录路径
    model: 大模型客户端
    
    Returns:
    tuple: (文件路径, 状态, 相对路径)
    """
    # 确定输出文件路径
    rel_path = os.path.relpath(transcript_file, input_dir)
    output_path = os.path.join(output_dir, rel_path)
    
    # 检查输出文件是否已存在且非空
    if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
        return (transcript_file, "skipped", rel_path)
    
    # 提取技巧
    result = distill_skill(transcript_file, output_path, model)
    status = "success" if result[1] else "failed"
    
    return (transcript_file, status, rel_path)

def process_transcripts(input_dir, output_dir, model_name="spark_deepseek_r1", n_jobs=10):
    """
    处理目录中的所有转写文本
    
    Parameters:
    input_dir (str): 转写文本目录
    output_dir (str): 结果输出目录
    model_name (str): 使用的模型名称
    n_jobs (int): 并行处理的作业数
    """
    # 加载模型
    print(f"正在加载模型 {model_name}...")
    model = get_by_name(model_name)
    print("模型加载完成。")
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 查找所有转写文本文件
    transcript_files = []
    
    for root, _, files in os.walk(input_dir):
        for file in files:
            if file.endswith('.txt'):
                file_path = os.path.join(root, file)
                transcript_files.append(file_path)
    
    print(f"找到 {len(transcript_files)} 个转写文本文件。")
    
    # 并行处理文件
    print(f"使用 {n_jobs} 个并行任务处理文件...")
    results = Parallel(n_jobs=n_jobs)(
        delayed(process_single_file)(transcript_file, input_dir, output_dir, model) 
        for transcript_file in tqdm(transcript_files, desc="提取技巧", unit="文件")
    )
    
    # 统计处理结果
    successful = sum(1 for _, status, _ in results if status == "success")
    skipped = sum(1 for _, status, _ in results if status == "skipped")
    failed = sum(1 for _, status, _ in results if status == "failed")
    
    # 打印处理信息
    for _, status, rel_path in results:
        if status == "success":
            print(f"已处理 {rel_path}")
        elif status == "skipped":
            print(f"跳过 {rel_path} (已存在提取结果)")
        elif status == "failed":
            print(f"处理失败 {rel_path}")
    
    # 打印摘要
    print("\n技巧提取摘要:")
    print(f"  总转写文件: {len(transcript_files)}")
    print(f"  成功提取: {successful}")
    print(f"  已跳过 (已存在): {skipped}")
    print(f"  失败: {failed}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="从转写文本中提取恋爱技巧")
    parser.add_argument("--input_dir", type=str, required=True, help="转写文本目录路径")
    parser.add_argument("--output_dir", type=str, required=True, help="结果输出目录路径")
    parser.add_argument("--model_name", type=str, default="spark_deepseek_r1", help="使用的模型名称")
    parser.add_argument("--n_jobs", type=int, default=10, help="并行处理的作业数")
    
    args = parser.parse_args()
    
    process_transcripts(
        args.input_dir,
        args.output_dir,
        args.model_name,
        args.n_jobs
    )